"""
PlumeSoft 量化代码库
作者：PlumeSoft
用途：为 optuna 设计的支持多进程、不依赖数据库、支持持久化保存的高性能炼丹基础框架
"""

import os
import time
import requests
import json
import numpy as np
import pandas as pd
from datetime import datetime
from filelock import FileLock
from traceback import format_exc


################################
# 调试相关函数
################################

def init_env(params: dict):
    set_log_file_name(params["study_log_file"], params["silent_mode"])
    set_profiler(params["use_profiler"])
    set_weixin_params(params["weixin_key"], params["weixin_prefix"], params["weixin_proxies"])


__study_log_file = None
__silent_mode = True


def set_log_file_name(path, silent_mode):
    global __study_log_file, __silent_mode
    __study_log_file = path
    __silent_mode = silent_mode
    os.makedirs(os.path.dirname(path), exist_ok=True)


def __write_file(path, type, *args):
    s_args = "\t".join([str(s) for s in args])
    with FileLock(path + ".lock"):
        with open(path, "a", encoding="utf8") as f:
            f.write(f"[{datetime.now()}][{type}] {s_args}\n")
            f.close()


# 打印带时间的日志
def log(*args, send_wx=False):
    print(f"[{datetime.now()}][LOG]", *args)
    if __study_log_file is not None:
        __write_file(__study_log_file, "LOG", *args)
    if send_wx:
        send_wx_message("\t".join([str(s) for s in args]))


# 打印用于调试的日志
def dbg(*args):
    if not __silent_mode:
        print(f"[{datetime.now()}][DBG]", *args)
    if __study_log_file is not None:
        __write_file(__study_log_file, "DBG", *args)


__weixin_key = None
__weixin_prefix = ""
__weixin_proxies = {}

def set_weixin_params(weixin_key, weixin_prefix, weixin_proxies):
    global __weixin_key, __weixin_prefix, __weixin_proxies
    __weixin_key = weixin_key
    __weixin_prefix = weixin_prefix
    __weixin_proxies = weixin_proxies


# 发送企业微信消息
def send_wx_message(content: str):
    try:
        if not __weixin_key:
            return

        local_time = time.localtime(time.time())
        s_time = time.strftime("%m-%d %H:%M:%S", local_time)
        s_prefix = "" if __weixin_prefix == "" else f"[{__weixin_prefix}]"
        msg = {
            "msgtype": "text",
            "text": {"content": f"[{s_time}]{s_prefix}\n{content.strip()}"},
        }

        headers = {"Content-Type": "application/json;charset=utf-8"}
        url = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=" + __weixin_key
        body = json.dumps(msg)
        requests.post(
            url, data=body, headers=headers, timeout=10, proxies=__weixin_proxies
        )

        return True
    except:
        return False


__profiler = None
__use_profiler = None


def set_profiler(name):
    global __use_profiler
    __use_profiler = name


def start_profiler(name):
    if __use_profiler == name:
        from pyinstrument import Profiler

        global __profiler
        __profiler = Profiler()
        __profiler.start()


def stop_profiler(name):
    if __use_profiler == name:
        __profiler.stop()
        log(__profiler.output_text(unicode=True, color=True))


################################
# 辅助函数
################################


def kill_process(pid=None):
    # 用于在 Ctrl+C 时强制退出所有进程
    try:
        if pid is None:
            pid = os.getpid()
        log(f"正在强制退出进程 {pid} 及所有子进程")
        import psutil

        p = psutil.Process(pid)
        for child in p.children(recursive=True):
            child.kill()
        p.kill()
    except OSError as e:
        log(e)


# 内存优化函数精简版本
# 该代码修改自：https://bbs.quantclass.cn/topic/index?id=10923
def reduce_mem_series_lite(col: pd.Series):
    col_type = col.dtypes
    if col_type == "int64":
        col = col.astype(np.int32, errors="ignore")
    elif col_type == "float64":
        col = col.astype(np.float32, errors="ignore")
    return col


def reduce_mem_df(df: pd.DataFrame, filters=[], dbg_info=False) -> pd.DataFrame:
    """
    使用更低精度的类型保存数据，优化内存占用
    """
    if df is None or df.empty:
        return df

    # 取得不在过滤列表中并且可以优化的字段
    col_list = [
        col
        for col in list(df.columns)
        if (col not in filters) and (df[col].dtypes in ["int64", "float64"])
    ]
    if len(col_list) == 0:
        return df

    # 用剩下的字段加优化后的字段组合新的 df
    if dbg_info:
        print(f"需要优化内存的字段: {col_list}")
    keep_list = [col for col in list(df.columns) if col not in col_list]
    df_list = [df[keep_list]] if len(keep_list) > 0 else []
    for col in col_list:
        df_list.append(reduce_mem_series_lite(df[col]))
    df = pd.concat(df_list, axis=1)
    return df


def merge_html(fig_list, output_folder, strategy_file, auto_open=True):
    # 创建合并后的网页文件
    os.makedirs(output_folder, exist_ok=True)
    merged_html_file = os.path.join(output_folder, f"{strategy_file}汇总.html")

    # 创建自定义HTML页面，嵌入fig对象的HTML内容
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
    <style>
        .figure-container {{
            display: flex;
            flex-direction: column;
            align-items: center;
        }}
    </style>
    </head>
    <body>"""
    for fig in fig_list:
        html_content += f"""
        <div class="figure-container">
            {fig}
        </div>
        """
    html_content += "</body> </html>"

    # 保存自定义HTML页面
    with open(merged_html_file, "w", encoding="utf-8") as f:
        f.write(html_content)

    if auto_open:
        res = os.system("start " + merged_html_file)
        if res != 0:
            os.system("open " + merged_html_file)


################################
# 安全读写数据文件相关函数
################################


def __try_load_dataframe(filename, fmt="pickle") -> pd.DataFrame:
    try:
        if os.path.isfile(filename):
            if fmt == "feather":
                return pd.read_feather(filename)
            else:
                return pd.read_pickle(filename)
        else:
            return None
    except:
        return None


def safe_get_file_size(filename):
    try:
        for fn in [filename, filename + ".tmp", filename + ".old"]:
            if os.path.isfile(fn):
                return os.path.getsize(fn)
        return None
    except:
        return None


def safe_isfile(filename) -> bool:
    try:
        for fn in [filename, filename + ".tmp", filename + ".old"]:
            if os.path.isfile(fn):
                return True
        return False
    except:
        return False


def safe_remove_file(filename) -> bool:
    try:
        for fn in [filename, filename + ".tmp", filename + ".old"]:
            if os.path.isfile(fn):
                os.remove(fn)
        return True
    except:
        return False


def safe_load_dataframe(filename, fmt="pickle") -> pd.DataFrame:
    for fn in [filename, filename + ".tmp", filename + ".old"]:
        data = __try_load_dataframe(fn, fmt)
        if not data is None:
            return data
    return None


def safe_load_feather_to_df(filename) -> pd.DataFrame:
    return safe_load_dataframe(filename, "feather")


def safe_save_dataframe(data: pd.DataFrame, filename, fmt="pickle") -> bool:
    try:
        # 先把已有文件改名 .old
        if os.path.isfile(filename):
            if os.path.isfile(filename + ".old"):
                os.remove(filename + ".old")
            os.rename(filename, filename + ".old")
        # 再保存为 .tmp 文件
        if fmt == "feather":
            data.to_feather(filename + ".tmp")
        else:
            data.to_pickle(filename + ".tmp")
        # 最后改名为原文件
        os.rename(filename + ".tmp", filename)
        # 删除旧文件
        if os.path.isfile(filename + ".old"):
            os.remove(filename + ".old")
        return True
    except:
        return False


def safe_save_df_to_feather(data: pd.DataFrame, filename) -> bool:
    return safe_save_dataframe(data, filename, "feather")


################################
# 代码生成相关函数
################################


def add_multi_code_to_stg_file(
    src_file, dst_file, code_list: list, encoding="utf8", match_all=True
) -> bool:
    # 增加代码到策略文件中
    try:
        with open(src_file, "r", encoding=encoding) as file_read:
            contents = file_read.read()
            succ_cnt = 0
            for flag1, flag2, code, append_mode in code_list:
                idx1 = contents.find(flag1)
                idx2 = contents.find(flag2)
                if idx1 > 0 and idx2 > 0:
                    if append_mode:
                        contents = contents.replace(flag2, code + flag2, 1)
                    else:
                        contents = (
                            contents[: idx1 + len(flag1)] + code + contents[idx2:]
                        )
                    succ_cnt += 1
            if succ_cnt == 0 or match_all and succ_cnt < len(code_list):
                log(
                    f"在源文件 {src_file} 中没有找到用于生成新文件的标志，请检查源文件内容。"
                )
                return False

        with open(dst_file, "w", encoding=encoding) as file_write:
            file_write.write(contents)
            return True
    except Exception as e:
        log(f"保存代码到策略文件 {dst_file} 发生异常:" + f"{str(e)}\n{format_exc()}")
        return False


def add_code_to_stg_file(
    src_file, dst_file, flag1, flag2, code, encoding="utf8", append_mode=True
) -> bool:
    # 增加代码到策略文件中
    return add_multi_code_to_stg_file(
        src_file, dst_file, [[flag1, flag2, code, append_mode]], encoding
    )
